#!/usr/bin/env python3
"""
高保真无人机导航系统模拟
包含完整的6自由度动力学模型、多传感器融合、改进的控制器和A*路径规划
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import math
import time
from queue import PriorityQueue
import logging
from logging.handlers import RotatingFileHandler
import csv
import os
from datetime import datetime
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial.transform import Rotation
import scipy.integrate

# 设置支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# 设置日志记录
def setup_logging():
    if not os.path.exists('logs'):
        os.makedirs('logs')

    log_filename = f"logs/drone_navigation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    log_format = '%(asctime)s - %(levelname)s - %(message)s'

    logging.basicConfig(
        level=logging.INFO,
        format=log_format,
        handlers=[
            RotatingFileHandler(log_filename, maxBytes=10 * 1024 * 1024, backupCount=5),
            logging.StreamHandler()
        ]
    )

    return log_filename


# 创建CSV数据记录器
class DataLogger:
    def __init__(self):
        if not os.path.exists('data'):
            os.makedirs('data')

        self.filename = f"data/drone_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"

        with open(self.filename, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([
                'timestamp', 'pos_x', 'pos_y', 'pos_z',
                'vel_x', 'vel_y', 'vel_z',
                'roll', 'pitch', 'yaw',
                'target_x', 'target_y', 'target_z',
                'error_x', 'error_y', 'error_z', 'total_error',
                'thrust', 'torque_x', 'torque_y', 'torque_z',
                'battery_remaining', 'motor_rpm_1', 'motor_rpm_2', 'motor_rpm_3', 'motor_rpm_4'
            ])

    def log_data(self, timestamp, state, target, control_input, battery, motor_rpms):
        position = state[0:3]
        error = target - position
        total_error = np.linalg.norm(error)

        with open(self.filename, 'a', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([
                timestamp,
                position[0], position[1], position[2],
                state[3], state[4], state[5],
                state[6], state[7], state[8],
                target[0], target[1], target[2],
                error[0], error[1], error[2], total_error,
                control_input[0], control_input[1], control_input[2], control_input[3],
                battery,
                motor_rpms[0], motor_rpms[1], motor_rpms[2], motor_rpms[3]
            ])


class PIDController:
    """PID控制器实现"""

    def __init__(self, kp, ki, kd, max_output=float('inf')):
        self.kp = np.array(kp)
        self.ki = np.array(ki)
        self.kd = np.array(kd)
        self.max_output = max_output

        self.integral = np.zeros(3)
        self.previous_error = np.zeros(3)
        self.previous_time = time.time()

        self.p_term = np.zeros(3)
        self.i_term = np.zeros(3)
        self.d_term = np.zeros(3)

    def compute(self, setpoint, current_value):
        current_time = time.time()
        dt = current_time - self.previous_time
        if dt <= 0:
            return np.zeros(3)

        error = setpoint - current_value

        # 比例项
        self.p_term = self.kp * error

        # 积分项 (带有抗饱和)
        self.integral += error * dt

        max_integral = np.zeros(3)
        for i in range(3):
            if self.ki[i] != 0:
                max_integral[i] = 5.0 / self.ki[i]
            else:
                max_integral[i] = float('inf')

        self.integral = np.clip(self.integral, -max_integral, max_integral)
        self.i_term = self.ki * self.integral

        # 微分项 (带有滤波)
        derivative = (error - self.previous_error) / dt
        self.d_term = self.kd * derivative

        # 计算总输出
        output = self.p_term + self.i_term + self.d_term

        # 限制输出
        output_norm = np.linalg.norm(output)
        if output_norm > self.max_output:
            output = output / output_norm * self.max_output

        # 更新状态
        self.previous_error = error
        self.previous_time = current_time

        logging.debug(f"PID - P: {self.p_term}, I: {self.i_term}, D: {self.d_term}, Output: {output}")

        return output


class MPCController:
    """改进版模型预测控制器"""

    def __init__(self, horizon=5, dt=0.1):
        self.horizon = horizon
        self.dt = dt

        # 权重矩阵
        self.Q = np.diag([10, 10, 10])
        self.R = np.diag([0.1, 0.1, 0.1])
        self.Qf = np.diag([20, 20, 20])

        # 系统矩阵
        self.A = np.eye(6)
        self.A[0:3, 3:6] = np.eye(3) * dt
        self.B = np.vstack([np.eye(3) * dt ** 2 / 2, np.eye(3) * dt])

        # 初始化
        self.last_control = np.zeros(3)
        self.last_error = np.zeros(3)

    def compute(self, current_state, reference_trajectory):
        # 简化的MPC实现
        current_pos = current_state[0:3]
        current_vel = current_state[3:6]

        error = reference_trajectory[0] - current_pos

        Kp = 3.0
        Kd = 2.0

        control = Kp * error + Kd * (self.last_error - error) / self.dt

        control_norm = np.linalg.norm(control)
        if control_norm > 5.0:
            control = control / control_norm * 5.0

        self.last_error = error
        self.last_control = control

        logging.debug(f"MPC - Error: {error}, Control: {control}")

        return control


class AStarPlanner:
    """改进版A*路径规划算法实现"""

    def __init__(self, grid_size=1.0):
        self.grid_size = grid_size

    def heuristic(self, a, b):
        return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2 + (a[2] - b[2]) ** 2)

    def plan(self, start, goal, obstacles):
        start_grid = (int(round(start[0] / self.grid_size)),
                      int(round(start[1] / self.grid_size)),
                      int(round(start[2] / self.grid_size)))
        goal_grid = (int(round(goal[0] / self.grid_size)),
                     int(round(goal[1] / self.grid_size)),
                     int(round(goal[2] / self.grid_size)))

        obstacle_grid = set()
        for obs in obstacles:
            obs_grid = (int(round(obs[0] / self.grid_size)),
                        int(round(obs[1] / self.grid_size)),
                        int(round(obs[2] / self.grid_size)))
            obstacle_grid.add(obs_grid)

        frontier = PriorityQueue()
        frontier.put((0, start_grid))
        came_from = {}
        cost_so_far = {}
        came_from[start_grid] = None
        cost_so_far[start_grid] = 0

        directions = []
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                for dz in [-1, 0, 1]:
                    if dx == 0 and dy == 0 and dz == 0:
                        continue
                    directions.append((dx, dy, dz))

        found_path = False
        while not frontier.empty():
            current = frontier.get()[1]

            if current == goal_grid:
                found_path = True
                break

            for dx, dy, dz in directions:
                next_node = (current[0] + dx, current[1] + dy, current[2] + dz)

                if next_node in obstacle_grid:
                    continue

                if abs(next_node[0]) > 50 or abs(next_node[1]) > 50 or next_node[2] < 0 or next_node[2] > 20:
                    continue

                new_cost = cost_so_far[current] + math.sqrt(dx * dx + dy * dy + dz * dz)
                if next_node not in cost_so_far or new_cost < cost_so_far[next_node]:
                    cost_so_far[next_node] = new_cost
                    priority = new_cost + self.heuristic(goal_grid, next_node)
                    frontier.put((priority, next_node))
                    came_from[next_node] = current

        if not found_path:
            logging.warning("A*路径规划失败，使用直线路径")
            return [start, goal]

        path = []
        current = goal_grid
        while current != start_grid:
            path.append(current)
            if current not in came_from:
                logging.error("A*路径重建失败")
                return [start, goal]
            current = came_from[current]
        path.append(start_grid)
        path.reverse()

        continuous_path = []
        for p in path:
            continuous_path.append(np.array([
                p[0] * self.grid_size,
                p[1] * self.grid_size,
                p[2] * self.grid_size
            ]))

        smoothed_path = self.smooth_path(continuous_path, obstacles)

        logging.info(f"A*规划完成，共{len(smoothed_path)}个航点")
        return smoothed_path

    def smooth_path(self, path, obstacles):
        if len(path) <= 2:
            return path

        smoothed = [path[0]]
        i = 0

        while i < len(path) - 1:
            j = len(path) - 1
            while j > i + 1:
                if self.is_line_clear(np.array(path[i]), np.array(path[j]), obstacles):
                    smoothed.append(path[j])
                    i = j
                    break
                j -= 1
            else:
                smoothed.append(path[i + 1])
                i += 1

        return smoothed

    def is_line_clear(self, start, end, obstacles):
        steps = 20
        for step in range(steps + 1):
            t = step / steps
            point = start + t * (end - start)

            for obs in obstacles:
                if np.linalg.norm(point - obs) < 1.5:
                    return False

        return True


class HighFidelityDrone:
    """高保真无人机模型 - 6自由度动力学"""

    def __init__(self, x=0, y=0, z=1, mass=1.0):
        # 物理参数
        self.mass = mass
        self.g = 9.81  # 重力加速度
        self.arm_length = 0.25  # 机臂长度
        self.max_thrust = 30.0  # 最大推力 (N)

        # 惯性矩阵 (假设对称结构)
        self.Ixx = 0.1
        self.Iyy = 0.1
        self.Izz = 0.2
        self.inertia = np.diag([self.Ixx, self.Iyy, self.Izz])
        self.inv_inertia = np.linalg.inv(self.inertia)

        # 空气动力学参数
        self.drag_coefficient = 0.1
        self.rotor_drag_coefficient = 0.01

        # 电机参数
        self.motor_time_constant = 0.02  # 电机响应时间常数
        self.motor_max_rpm = 10000  # 最大转速 (RPM)
        self.motor_thrust_coefficient = 8.54858e-6  # 推力系数
        self.motor_torque_coefficient = 1.0e-7  # 扭矩系数

        # 初始状态 [x, y, z, vx, vy, vz, φ, θ, ψ, p, q, r]
        self.state = np.zeros(12)
        self.state[0:3] = [x, y, z]  # 初始位置
        self.state[6:9] = [0, 0, 0]  # 初始姿态 (滚转, 俯仰, 偏航)

        # 电机状态
        self.motor_rpms = np.zeros(4)  # 四个电机的转速 (RPM)

        # 电池参数
        self.battery_capacity = 2200  # mAh
        self.battery_voltage = 11.1  # V
        self.battery_remaining = 1.0  # 剩余电量百分比
        self.power_consumption = 0.0  # 当前功率消耗 (W)

        # 控制器
        self.position_controller = PIDController(
            kp=np.array([2.0, 2.0, 6.0]),
            ki=np.array([0.1, 0.1, 0.1]),
            kd=np.array([1.0, 1.0, 3.0]),
            max_output=25.0
        )
        self.attitude_controller = PIDController(
            kp=np.array([8.0, 8.0, 6.0]),
            ki=np.array([0.5, 0.5, 0.5]),
            kd=np.array([0.2, 0.2, 0.2]),
            max_output=5.0
        )

        # 历史数据记录
        self.trajectory = []
        self.control_inputs = []
        self.motor_history = []
        self.battery_history = []

        # 步数计数器
        self.step_count = 0

        # 环境参数
        self.wind = np.zeros(3)  # 风速向量

        logging.info("高保真无人机模型初始化完成")

    def dynamics(self, t, state, control_input):
        """6自由度动力学方程"""
        # 解包状态变量
        x, y, z, vx, vy, vz, phi, theta, psi, p, q, r = state
        thrust, tau_phi, tau_theta, tau_psi = control_input

        # 旋转矩阵 (从机体坐标系到世界坐标系)
        R = Rotation.from_euler('zyx', [psi, theta, phi], degrees=False).as_matrix()

        # 重力向量
        gravity = np.array([0, 0, self.g])

        # 总推力 (在机体坐标系中，推力沿z轴负方向)
        thrust_body = np.array([0, 0, -thrust])

        # 转换到世界坐标系
        thrust_world = R @ thrust_body

        # 空气阻力 (与速度平方成正比，方向相反)
        velocity = np.array([vx, vy, vz])
        drag_force = -self.drag_coefficient * np.linalg.norm(velocity) * velocity

        # 线加速度
        acceleration = thrust_world / self.mass + gravity + drag_force / self.mass

        # 角动力学
        # 角速度向量
        omega = np.array([p, q, r])

        # 角加速度 (欧拉方程)
        torque = np.array([tau_phi, tau_theta, tau_psi])
        omega_dot = self.inv_inertia @ (torque - np.cross(omega, self.inertia @ omega))

        # 姿态角变化率 (从角速度到欧拉角变化率)
        euler_rates = np.array([
            [1, np.sin(phi) * np.tan(theta), np.cos(phi) * np.tan(theta)],
            [0, np.cos(phi), -np.sin(phi)],
            [0, np.sin(phi) / np.cos(theta), np.cos(phi) / np.cos(theta)]
        ]) @ omega

        # 状态导数
        state_dot = np.zeros(12)
        state_dot[0:3] = velocity  # 位置变化率
        state_dot[3:6] = acceleration  # 速度变化率
        state_dot[6:9] = euler_rates  # 欧拉角变化率
        state_dot[9:12] = omega_dot  # 角加速度

        return state_dot

    def motor_dynamics(self, desired_rpms, dt):
        """电机动力学模型"""
        # 一阶响应模型
        self.motor_rpms += (desired_rpms - self.motor_rpms) * (1 - np.exp(-dt / self.motor_time_constant))

        # 限制电机转速
        self.motor_rpms = np.clip(self.motor_rpms, 0, self.motor_max_rpm)

        return self.motor_rpms

    def control_allocation(self, thrust, torques):
        """控制分配 - 将总推力和力矩分配到四个电机"""
        # 混控矩阵 (X型布局)
        mix_matrix = np.array([
            [1, 1, 1, 1],  # 总推力
            [0, -1, 0, 1],  # 滚转力矩
            [-1, 0, 1, 0],  # 俯仰力矩
            [1, -1, 1, -1]  # 偏航力矩
        ])

        # 计算电机指令 (伪逆求解)
        motor_commands = np.linalg.pinv(mix_matrix) @ np.array([thrust, torques[0], torques[1], torques[2]])

        # 将电机指令转换为RPM
        # 推力 = k_thrust * RPM²
        desired_rpms = np.sqrt(np.maximum(motor_commands, 0) / self.motor_thrust_coefficient)

        return desired_rpms

    def update_battery(self, motor_rpms, dt):
        """更新电池状态"""
        # 计算电机功率消耗 (功率 ≈ k * RPM³)
        power_per_motor = 1e-9 * motor_rpms ** 3  # 简化模型
        total_power = np.sum(power_per_motor)

        # 更新电池电量
        energy_consumed = total_power * dt / 3600  # 转换为Ah
        self.battery_remaining -= energy_consumed / (self.battery_capacity / 1000)  # 转换为百分比

        # 限制电池电量
        self.battery_remaining = max(0, self.battery_remaining)

        self.power_consumption = total_power

        return self.battery_remaining

    def update(self, target_position, dt=0.1):
        """更新无人机状态"""
        self.step_count += 1

        # 位置控制 (外环)
        position_error = target_position - self.state[0:3]
        desired_acceleration = self.position_controller.compute(position_error, self.state[3:6])

        # 考虑重力和当前姿态
        gravity = np.array([0, 0, self.g])
        R = Rotation.from_euler('zyx', self.state[6:9], degrees=False).as_matrix()
        current_thrust_direction = R @ np.array([0, 0, 1])

        # 计算所需总推力
        desired_thrust_magnitude = self.mass * (np.linalg.norm(desired_acceleration - gravity) + 0.1)
        desired_thrust_magnitude = np.clip(desired_thrust_magnitude, 0, self.max_thrust)

        # 计算期望姿态 (使推力方向对齐期望加速度方向)
        # 简化处理：假设偏航角保持不变
        desired_thrust_direction = (desired_acceleration - gravity)
        desired_thrust_direction = desired_thrust_direction / np.linalg.norm(desired_thrust_direction)

        # 计算期望的滚转和俯仰角
        desired_roll = np.arcsin(desired_thrust_direction[1])
        desired_pitch = np.arctan2(-desired_thrust_direction[0], -desired_thrust_direction[2])
        desired_yaw = self.state[8]  # 保持当前偏航角

        # 姿态控制 (内环)
        desired_attitude = np.array([desired_roll, desired_pitch, desired_yaw])
        attitude_error = desired_attitude - self.state[6:9]

        # 角速率控制
        desired_rates = self.attitude_controller.compute(attitude_error, self.state[9:12])

        # 计算控制力矩
        # 简化处理：直接使用PD控制
        torques = np.array([
            2.0 * attitude_error[0] + 0.5 * desired_rates[0],  # 滚转力矩
            2.0 * attitude_error[1] + 0.5 * desired_rates[1],  # 俯仰力矩
            1.5 * attitude_error[2] + 0.3 * desired_rates[2]  # 偏航力矩
        ])

        # 控制分配
        desired_rpms = self.control_allocation(desired_thrust_magnitude, torques)

        # 电机动力学
        motor_rpms = self.motor_dynamics(desired_rpms, dt)

        # 计算实际产生的推力和力矩
        actual_thrust = np.sum(self.motor_thrust_coefficient * motor_rpms ** 2)
        actual_torques = np.array([
            self.arm_length * self.motor_thrust_coefficient * (motor_rpms[3] ** 2 - motor_rpms[1] ** 2),
            self.arm_length * self.motor_thrust_coefficient * (motor_rpms[2] ** 2 - motor_rpms[0] ** 2),
            self.motor_torque_coefficient * (
                        motor_rpms[0] ** 2 - motor_rpms[1] ** 2 + motor_rpms[2] ** 2 - motor_rpms[3] ** 2)
        ])

        # 更新电池状态
        battery_remaining = self.update_battery(motor_rpms, dt)

        # 积分动力学方程
        control_input = np.array([actual_thrust, actual_torques[0], actual_torques[1], actual_torques[2]])
        sol = scipy.integrate.solve_ivp(
            fun=lambda t, y: self.dynamics(t, y, control_input),
            t_span=[0, dt],
            y0=self.state,
            t_eval=[dt],
            method='RK45',
            rtol=1e-6,
            atol=1e-9
        )

        # 更新状态
        self.state = sol.y[:, -1]

        # 记录数据
        self.trajectory.append(self.state[0:3].copy())
        self.control_inputs.append(control_input.copy())
        self.motor_history.append(motor_rpms.copy())
        self.battery_history.append(battery_remaining)

        # 每10步记录一次详细信息
        if self.step_count % 10 == 0:
            error = target_position - self.state[0:3]
            logging.info(
                f"步数 {self.step_count}: 位置={self.state[0:3]}, 目标={target_position}, "
                f"误差={error}, 推力={actual_thrust:.2f}N, 电池={battery_remaining * 100:.1f}%"
            )

        # 检查是否到达目标
        distance = np.linalg.norm(target_position - self.state[0:3])
        return distance < 0.5

    def get_state(self):
        return self.state.copy()

    def get_trajectory(self):
        return np.array(self.trajectory)

    def get_control_inputs(self):
        return np.array(self.control_inputs)

    def get_motor_history(self):
        return np.array(self.motor_history)

    def get_battery_history(self):
        return np.array(self.battery_history)


class NavigationSystem:
    def __init__(self):
        # 设置日志
        self.log_filename = setup_logging()
        self.data_logger = DataLogger()

        logging.info("初始化高保真导航系统...")

        # 创建障碍物
        self.obstacles = [
            np.array([5, 5, 2]),
            np.array([8, 3, 3]),
            np.array([12, 8, 4]),
            np.array([15, 12, 5]),
            np.array([7, 10, 3]),
            np.array([10, 15, 4]),
            np.array([13, 5, 3]),
            np.array([18, 8, 4])
        ]

        # 起点和终点
        self.start = np.array([0, 0, 1])
        self.goal = np.array([20, 15, 6])

        # 路径规划器
        self.planner = AStarPlanner(grid_size=1.0)

        # 规划路径
        self.waypoints = self.plan_path()

        # 当前目标航点索引
        self.current_waypoint_idx = 1

        # 创建高保真无人机实例
        self.drone = HighFidelityDrone(
            x=self.start[0],
            y=self.start[1],
            z=self.start[2]
        )

        # 是否完成所有航点
        self.mission_complete = False

        # 开始时间
        self.start_time = time.time()

        logging.info(f"高保真导航系统初始化完成，共{len(self.waypoints)}个航点")

    def plan_path(self):
        """使用A*算法规划路径"""
        logging.info("正在规划路径...")
        path = self.planner.plan(self.start, self.goal, self.obstacles)

        if not path:
            logging.warning("无法找到路径! 使用直线路径")
            return [self.start, self.goal]

        logging.info(f"路径规划完成，共{len(path)}个航点")
        return path

    def update(self, dt=0.1):
        """更新导航系统状态"""
        if self.mission_complete:
            return True

        # 获取当前目标航点
        target = self.waypoints[self.current_waypoint_idx]

        # 更新无人机位置
        reached = self.drone.update(target, dt)

        # 记录数据
        current_time = time.time() - self.start_time
        motor_rpms = self.drone.motor_rpms
        self.data_logger.log_data(
            current_time,
            self.drone.get_state(),
            target,
            self.drone.control_inputs[-1] if self.drone.control_inputs else np.zeros(4),
            self.drone.battery_remaining,
            motor_rpms
        )

        # 如果到达当前航点，则转向下一个航点
        if reached:
            logging.info(f"到达航点 {self.current_waypoint_idx}: {target}")
            self.current_waypoint_idx += 1

            # 检查是否完成所有航点
            if self.current_waypoint_idx >= len(self.waypoints):
                logging.info("任务完成!")
                self.mission_complete = True
                return True

        # 检查电池电量
        if self.drone.battery_remaining <= 0.1:  # 电量低于10%
            logging.warning("电池电量不足，启动紧急降落程序")
            self.emergency_landing()
            self.mission_complete = True
            return True

        return False

    def emergency_landing(self):
        """紧急降落程序"""
        logging.info("执行紧急降落程序")
        # 简化处理：直接下降到地面
        target = self.drone.state[0:3].copy()
        target[2] = 0  # 设置高度为0

        # 缓慢下降直到触地
        while self.drone.state[2] > 0.1:
            self.drone.update(target, 0.1)
            time.sleep(0.1)

        logging.info("紧急降落完成")

    def get_drone_position(self):
        return self.drone.get_state()[0:3]

    def get_waypoints(self):
        return self.waypoints

    def get_trajectory(self):
        return self.drone.get_trajectory()

    def get_control_inputs(self):
        return self.drone.get_control_inputs()

    def get_obstacles(self):
        return self.obstacles

    def print_summary(self):
        """打印任务摘要"""
        trajectory = self.get_trajectory()
        total_distance = 0
        for i in range(1, len(trajectory)):
            total_distance += np.linalg.norm(trajectory[i] - trajectory[i - 1])

        energy_consumed = (1 - self.drone.battery_remaining) * self.drone.battery_capacity / 1000  # Ah
        total_energy = energy_consumed * self.drone.battery_voltage  # Wh

        logging.info(f"任务摘要: 总飞行距离={total_distance:.2f}m, 航点数={len(self.waypoints)}")
        logging.info(f"能量消耗: {total_energy:.2f}Wh, 电池剩余={self.drone.battery_remaining * 100:.1f}%")
        logging.info(f"日志文件: {self.log_filename}")


class Simulation:
    def __init__(self, controller_type="PID"):
        # 创建导航系统
        self.nav_system = NavigationSystem()

        # 创建图形
        self.fig = plt.figure(figsize=(20, 12))

        # 创建3D子图
        self.ax_3d = self.fig.add_subplot(241, projection='3d')
        self.ax_position = self.fig.add_subplot(242)
        self.ax_attitude = self.fig.add_subplot(243)
        self.ax_control = self.fig.add_subplot(244)
        self.ax_motors = self.fig.add_subplot(245)
        self.ax_battery = self.fig.add_subplot(246)
        self.ax_energy = self.fig.add_subplot(247)
        self.ax_error = self.fig.add_subplot(248)

        # 设置图表标题和标签
        self.ax_3d.set_title('3D轨迹')
        self.ax_3d.set_xlabel('X (m)')
        self.ax_3d.set_ylabel('Y (m)')
        self.ax_3d.set_zlabel('Z (m)')

        # 设置坐标轴范围
        self.ax_3d.set_xlim(-2, 25)
        self.ax_3d.set_ylim(-2, 20)
        self.ax_3d.set_zlim(0, 10)

        # 绘制起点和终点
        start = self.nav_system.start
        goal = self.nav_system.goal
        self.ax_3d.scatter(start[0], start[1], start[2], c='green', marker='o', s=100, label='起点')
        self.ax_3d.scatter(goal[0], goal[1], goal[2], c='red', marker='o', s=100, label='终点')

        # 绘制航点
        waypoints = self.nav_system.get_waypoints()
        wp_x = [wp[0] for wp in waypoints]
        wp_y = [wp[1] for wp in waypoints]
        wp_z = [wp[2] for wp in waypoints]
        self.ax_3d.plot(wp_x, wp_y, wp_z, 'r--', alpha=0.5, label='规划路径')

        # 绘制障碍物
        obstacles = self.nav_system.get_obstacles()
        for obs in obstacles:
            self.draw_cube(obs, 1.0, 'gray')

        # 初始化无人机位置标记
        self.drone_marker, = self.ax_3d.plot([], [], [], 'bo', markersize=8, label='无人机')

        # 初始化轨迹线
        self.trajectory_line, = self.ax_3d.plot([], [], [], 'b-', alpha=0.7, linewidth=2, label='实际轨迹')

        # 初始化当前目标标记
        self.target_marker, = self.ax_3d.plot([], [], [], 'rx', markersize=10, label='当前目标')

        # 添加图例
        self.ax_3d.legend()

        # 设置视角
        self.ax_3d.view_init(elev=20, azim=45)

        # 初始化其他图表
        self.init_position_plot()
        self.init_attitude_plot()
        self.init_control_plot()
        self.init_motor_plot()
        self.init_battery_plot()
        self.init_energy_plot()
        self.init_error_plot()

        # 动画对象
        self.animation = None

        # 帧计数器
        self.frame_count = 0

        # 数据记录
        self.time_history = []
        self.position_history = []
        self.attitude_history = []
        self.control_history = []
        self.motor_history = []
        self.battery_history = []
        self.error_history = []

    def draw_cube(self, center, size, color):
        """绘制立方体障碍物"""
        r = size / 2.0
        vertices = [
            [center[0] - r, center[1] - r, center[2] - r],
            [center[0] + r, center[1] - r, center[2] - r],
            [center[0] + r, center[1] + r, center[2] - r],
            [center[0] - r, center[1] + r, center[2] - r],
            [center[0] - r, center[1] - r, center[2] + r],
            [center[0] + r, center[1] - r, center[2] + r],
            [center[0] + r, center[1] + r, center[2] + r],
            [center[0] - r, center[1] + r, center[2] + r]
        ]

        faces = [
            [0, 1, 2, 3],
            [4, 5, 6, 7],
            [0, 1, 5, 4],
            [2, 3, 7, 6],
            [0, 3, 7, 4],
            [1, 2, 6, 5]
        ]

        poly_verts = []
        for face in faces:
            poly_verts.append([vertices[i] for i in face])

        collection = Poly3DCollection(poly_verts, facecolors=color, linewidths=0.5, edgecolors='k', alpha=0.25)
        self.ax_3d.add_collection3d(collection)

    def init_position_plot(self):
        """初始化位置图表"""
        self.ax_position.set_title('位置跟踪')
        self.ax_position.set_xlabel('时间 (s)')
        self.ax_position.set_ylabel('位置 (m)')
        self.ax_position.grid(True)

        self.position_x_line, = self.ax_position.plot([], [], 'r-', label='X')
        self.position_y_line, = self.ax_position.plot([], [], 'g-', label='Y')
        self.position_z_line, = self.ax_position.plot([], [], 'b-', label='Z')
        self.ax_position.legend()

    def init_attitude_plot(self):
        """初始化姿态图表"""
        self.ax_attitude.set_title('姿态角')
        self.ax_attitude.set_xlabel('时间 (s)')
        self.ax_attitude.set_ylabel('角度 (rad)')
        self.ax_attitude.grid(True)

        self.attitude_roll_line, = self.ax_attitude.plot([], [], 'r-', label='Roll')
        self.attitude_pitch_line, = self.ax_attitude.plot([], [], 'g-', label='Pitch')
        self.attitude_yaw_line, = self.ax_attitude.plot([], [], 'b-', label='Yaw')
        self.ax_attitude.legend()

    def init_control_plot(self):
        """初始化控制输入图表"""
        self.ax_control.set_title('控制输入')
        self.ax_control.set_xlabel('时间 (s)')
        self.ax_control.set_ylabel('输入值')
        self.ax_control.grid(True)

        self.control_thrust_line, = self.ax_control.plot([], [], 'r-', label='推力')
        self.control_tau_x_line, = self.ax_control.plot([], [], 'g-', label='τx')
        self.control_tau_y_line, = self.ax_control.plot([], [], 'b-', label='τy')
        self.control_tau_z_line, = self.ax_control.plot([], [], 'm-', label='τz')
        self.ax_control.legend()

    def init_motor_plot(self):
        """初始化电机转速图表"""
        self.ax_motors.set_title('电机转速')
        self.ax_motors.set_xlabel('时间 (s)')
        self.ax_motors.set_ylabel('RPM')
        self.ax_motors.grid(True)

        self.motor_1_line, = self.ax_motors.plot([], [], 'r-', label='电机1')
        self.motor_2_line, = self.ax_motors.plot([], [], 'g-', label='电机2')
        self.motor_3_line, = self.ax_motors.plot([], [], 'b-', label='电机3')
        self.motor_4_line, = self.ax_motors.plot([], [], 'm-', label='电机4')
        self.ax_motors.legend()

    def init_battery_plot(self):
        """初始化电池图表"""
        self.ax_battery.set_title('电池状态')
        self.ax_battery.set_xlabel('时间 (s)')
        self.ax_battery.set_ylabel('剩余电量 (%)')
        self.ax_battery.grid(True)

        self.battery_line, = self.ax_battery.plot([], [], 'b-')
        self.ax_battery.set_ylim(0, 100)

    def init_energy_plot(self):
        """初始化能量消耗图表"""
        self.ax_energy.set_title('能量消耗')
        self.ax_energy.set_xlabel('时间 (s)')
        self.ax_energy.set_ylabel('功率 (W)')
        self.ax_energy.grid(True)

        self.energy_line, = self.ax_energy.plot([], [], 'r-')

    def init_error_plot(self):
        """初始化误差图表"""
        self.ax_error.set_title('跟踪误差')
        self.ax_error.set_xlabel('时间 (s)')
        self.ax_error.set_ylabel('误差 (m)')
        self.ax_error.grid(True)

        self.error_x_line, = self.ax_error.plot([], [], 'r-', label='X误差')
        self.error_y_line, = self.ax_error.plot([], [], 'g-', label='Y误差')
        self.error_z_line, = self.ax_error.plot([], [], 'b-', label='Z误差')
        self.error_total_line, = self.ax_error.plot([], [], 'k--', label='总误差')
        self.ax_error.legend()

    def update_plot(self, frame):
        """更新动画帧"""
        self.frame_count += 1

        # 更新导航系统
        complete = self.nav_system.update(dt=0.1)

        # 获取无人机当前位置和状态
        drone_state = self.nav_system.drone.get_state()
        drone_pos = drone_state[0:3]

        # 获取当前目标
        if self.nav_system.current_waypoint_idx < len(self.nav_system.waypoints):
            target_pos = self.nav_system.waypoints[self.nav_system.current_waypoint_idx]
        else:
            target_pos = self.nav_system.waypoints[-1]

        # 更新无人机位置标记
        self.drone_marker.set_data([drone_pos[0]], [drone_pos[1]])
        self.drone_marker.set_3d_properties([drone_pos[2]])

        # 更新目标位置标记
        self.target_marker.set_data([target_pos[0]], [target_pos[1]])
        self.target_marker.set_3d_properties([target_pos[2]])

        # 更新轨迹线
        trajectory = self.nav_system.get_trajectory()
        if len(trajectory) > 1:
            self.trajectory_line.set_data(trajectory[:, 0], trajectory[:, 1])
            self.trajectory_line.set_3d_properties(trajectory[:, 2])

        # 记录数据用于绘图
        current_time = time.time() - self.nav_system.start_time
        self.time_history.append(current_time)
        self.position_history.append(drone_pos)
        self.attitude_history.append(drone_state[6:9])

        if self.nav_system.drone.control_inputs:
            self.control_history.append(self.nav_system.drone.control_inputs[-1])

        if len(self.nav_system.drone.motor_history) > 0:
            self.motor_history.append(self.nav_system.drone.motor_history[-1])

        self.battery_history.append(self.nav_system.drone.battery_remaining * 100)

        error = target_pos - drone_pos
        self.error_history.append(np.concatenate([error, [np.linalg.norm(error)]]))

        # 更新所有图表
        if len(self.time_history) > 1:
            time_array = np.array(self.time_history)

            # 更新位置图表
            position_array = np.array(self.position_history)
            self.position_x_line.set_data(time_array, position_array[:, 0])
            self.position_y_line.set_data(time_array, position_array[:, 1])
            self.position_z_line.set_data(time_array, position_array[:, 2])
            self.ax_position.relim()
            self.ax_position.autoscale_view()

            # 更新姿态图表
            attitude_array = np.array(self.attitude_history)
            self.attitude_roll_line.set_data(time_array, attitude_array[:, 0])
            self.attitude_pitch_line.set_data(time_array, attitude_array[:, 1])
            self.attitude_yaw_line.set_data(time_array, attitude_array[:, 2])
            self.ax_attitude.relim()
            self.ax_attitude.autoscale_view()

            # 更新控制输入图表
            if len(self.control_history) > 0:
                control_array = np.array(self.control_history)
                self.control_thrust_line.set_data(time_array[-len(control_array):], control_array[:, 0])
                self.control_tau_x_line.set_data(time_array[-len(control_array):], control_array[:, 1])
                self.control_tau_y_line.set_data(time_array[-len(control_array):], control_array[:, 2])
                self.control_tau_z_line.set_data(time_array[-len(control_array):], control_array[:, 3])
                self.ax_control.relim()
                self.ax_control.autoscale_view()

            # 更新电机转速图表
            if len(self.motor_history) > 0:
                motor_array = np.array(self.motor_history)
                self.motor_1_line.set_data(time_array[-len(motor_array):], motor_array[:, 0])
                self.motor_2_line.set_data(time_array[-len(motor_array):], motor_array[:, 1])
                self.motor_3_line.set_data(time_array[-len(motor_array):], motor_array[:, 2])
                self.motor_4_line.set_data(time_array[-len(motor_array):], motor_array[:, 3])
                self.ax_motors.relim()
                self.ax_motors.autoscale_view()

            # 更新电池图表
            battery_array = np.array(self.battery_history)
            self.battery_line.set_data(time_array, battery_array)
            self.ax_battery.relim()
            self.ax_battery.autoscale_view()

            # 更新能量图表 (简化处理)
            if hasattr(self.nav_system.drone, 'power_consumption'):
                power_array = np.full_like(time_array, self.nav_system.drone.power_consumption)
                self.energy_line.set_data(time_array, power_array)
                self.ax_energy.relim()
                self.ax_energy.autoscale_view()

            # 更新误差图表
            error_array = np.array(self.error_history)
            self.error_x_line.set_data(time_array, error_array[:, 0])
            self.error_y_line.set_data(time_array, error_array[:, 1])
            self.error_z_line.set_data(time_array, error_array[:, 2])
            self.error_total_line.set_data(time_array, error_array[:, 3])
            self.ax_error.relim()
            self.ax_error.autoscale_view()

        # 如果任务完成，停止动画
        if complete:
            self.animation.event_source.stop()
            self.nav_system.print_summary()
            logging.info("模拟完成")

        return (self.drone_marker, self.trajectory_line, self.target_marker,
                self.position_x_line, self.position_y_line, self.position_z_line,
                self.attitude_roll_line, self.attitude_pitch_line, self.attitude_yaw_line,
                self.control_thrust_line, self.control_tau_x_line, self.control_tau_y_line, self.control_tau_z_line,
                self.motor_1_line, self.motor_2_line, self.motor_3_line, self.motor_4_line,
                self.battery_line, self.energy_line,
                self.error_x_line, self.error_y_line, self.error_z_line, self.error_total_line)

    def run(self):
        """运行模拟"""
        logging.info("开始高保真无人机导航模拟...")
        logging.info(f"起点: {self.nav_system.start}")
        logging.info(f"终点: {self.nav_system.goal}")
        logging.info(f"障碍物数量: {len(self.nav_system.get_obstacles())}")

        # 创建动画
        self.animation = FuncAnimation(
            self.fig, self.update_plot, frames=1000,
            interval=50, blit=False, repeat=False
        )

        plt.tight_layout()
        plt.show()


# 运行模拟
if __name__ == "__main__":
    sim = Simulation()
    sim.run()